from pulp import LpMaximize, LpProblem, LpVariable, lpSum, PULP_CBC_CMD
from common_import import *


def tianabc():
    year_data = {}
    years = list(range(2024, 2031))
    cumulative_growth_wheat_corn = 1.0
    cumulative_sales_variation = 1.0
    cumulative_yield_variation = 1.0
    cumulative_price_decrease = 1.0
    yangdu = 1.0
    cost = 1.0
    for year in years:
        cumulative_growth_wheat_corn *= 1.075
        cumulative_sales_variation *= 1
        cumulative_yield_variation *= 1
        cumulative_price_decrease *= 0.97
        yangdu *= 0.95
        cost *= 1.05
        year_data[year] = {
            "玉米销量": round(cumulative_growth_wheat_corn, 4),
            "其他销量": round(cumulative_sales_variation, 4),
            "亩产量": round(cumulative_yield_variation, 4),
            "食用菌价格": round(cumulative_price_decrease, 4),
            "羊肚菌价格": round(yangdu, 4),
            "种植成本": round(cost, 4),
        }
    samples_list = year_data
    # 创建问题
    model = LpProblem(name="maximize-profit", sense=LpMaximize)

    # 地块、作物、时间范围
    lands = [
        "A1",
        "A2",
        "A3",
        "A4",
        "A5",
        "A6",
        "B1",
        "B2",
        "B3",
        "B4",
        "B5",
        "B6",
        "B7",
        "B8",
        "B9",
        "B10",
        "B11",
        "B12",
        "B13",
        "B14",
        "C1",
        "C2",
        "C3",
        "C4",
        "C5",
        "C6",
    ]
    crops = range(1, 16)  # 作物编号
    years = range(2024, 2031)  # 时间范围：2024到2030

    # 面积、产量等数据
    area_of_land = {
        "A1": 80,
        "A2": 55,
        "A3": 35,
        "A4": 72,
        "A5": 68,
        "A6": 55,
        "B1": 60,
        "B2": 46,
        "B3": 40,
        "B4": 28,
        "B5": 25,
        "B6": 86,
        "B7": 55,
        "B8": 44,
        "B9": 50,
        "B10": 25,
        "B11": 60,
        "B12": 45,
        "B13": 35,
        "B14": 20,
        "C1": 15,
        "C2": 13,
        "C3": 15,
        "C4": 18,
        "C5": 27,
        "C6": 20,
    }

    total_sales = {
        1: 57000,
        2: 21850,
        3: 22400,
        4: 33040,
        5: 9875,
        6: 170840,
        7: 132750,
        8: 71400,
        9: 30000,
        10: 12500,
        11: 1500,
        12: 35100,
        13: 36000,
        14: 14000,
        15: 10000,
    }
    x = LpVariable.dicts(
        "x",
        ((i, j, t) for i in lands for j in crops for t in years),
        lowBound=0,
        cat="Inte",
    )
    z = LpVariable.dicts(
        "z", ((i, j, t) for i in lands for j in crops for t in years), cat="Binary"
    )

    # 新增变量：表示每种作物每年超出的部分
    overproduction = LpVariable.dicts(
        "overproduction",
        ((j, t) for j in crops for t in years),
        lowBound=0,
        cat="Continuous",
    )

    # 目标函数：最大化总利润，并惩罚超出部分
    model += lpSum(
        mapping.get_profit_idtype(j, mapping.get_land_type(i), "单季")
        * x[i, j, t]
        * mapping.get_yield_idtype(j, mapping.get_land_type(i), "单季")
        * samples_list[t]["亩产量"]
        - mapping.get_cost_idtype(j, mapping.get_land_type(i), "单季")
        * (samples_list[t]["种植成本"] - 1)
        for i in lands
        for j in crops
        for t in years
    ) - lpSum(
        overproduction[j, t] * mapping.get_priceperjin_idtype(j, "单季")
        for j in crops
        for t in years
    )

    # 约束条件1: 每个作物在所有地的一年总产量小于等于总销量 + 超出的部分
    for j in crops:
        for t in years:
            total_yield = lpSum(
                mapping.get_yield_idtype(j, mapping.get_land_type(i), "单季")
                * samples_list[t]["亩产量"]
                * x[i, j, t]
                for i in lands
            )
            if j == 6 or j == 7:
                model += (
                    total_yield
                    <= total_sales[j] * samples_list[t]["玉米销量"]
                    + overproduction[j, t]
                )  # 允许超出部分
                model += (
                    overproduction[j, t]
                    >= total_yield - total_sales[j] * samples_list[t]["玉米销量"]
                )  # 计算超出的部分
            else:
                model += (
                    total_yield
                    <= total_sales[j] * samples_list[t]["其他销量"]
                    + overproduction[j, t]
                )  # 允许超出部分
                model += (
                    overproduction[j, t]
                    >= total_yield - total_sales[j] * samples_list[t]["其他销量"]
                )  # 计算超出的部分

    # 约束条件: 2023年的已知种植方案影响到2024年
    known_planting_2023 = {
        "A1": 6,
        "A2": 7,
        "A3": 7,
        "A4": 1,
        "A5": 4,
        "A6": 8,
        "B1": 6,
        "B2": 2,
        "B3": 3,
        "B4": 4,
        "B5": 5,
        "B6": 8,
        "B7": 6,
        "B8": 8,
        "B9": 9,
        "B10": 10,
        "B11": 1,
        "B12": 7,
        "B13": 14,
        "B14": 15,
        "C1": 11,
        "C2": 12,
        "C3": 1,
        "C4": 13,
        "C5": 6,
        "C6": 3,
    }

    # 约束条件2: 每个地任意连续三年种植的豆类粮食面积大于等于该地的面积
    bean_crops = [1, 2, 3, 4, 5]
    for i in lands:
        # 考虑2023, 2024, 2025三个连续年份
        model += (
            lpSum(
                x[i, j, 2024]
                + x[i, j, 2025]
                + (100 if (i, j, 2023) in known_planting_2023 else 0)
                for j in bean_crops
            )
            >= area_of_land[i]
        )

        # 其他年份的连续三年
        for t in range(2024, 2029):
            model += (
                lpSum(x[i, j, t + k] for j in bean_crops for k in range(3))
                >= area_of_land[i]
            )
        pass

    # 约束条件3: 每个地每年种植的作物总面积小于等于该地的面积
    for i in lands:
        for t in years:
            model += lpSum(x[i, j, t] for j in crops) <= area_of_land[i]

    # 约束条件4: 不能重茬种植
    M = 100
    for i in lands:
        for j in crops:
            for t in range(2024, 2030):
                model += x[i, j, t] <= M * z[i, j, t]
                model += z[i, j, t] + z[i, j, t + 1] <= 1

    # 2024不能重茬种植
    for i, j in known_planting_2023.items():
        model += z[i, j, 2024] == 0  # 2024年不能重茬种植
    # 启用CBC求解器的详细输出
    # solver = PULP_CBC_CMD(msg=False, gapRel=0.05)
    solver = PULP_CBC_CMD(msg=False, gapRel=0.05, timeLimit=10)

    # 求解模型
    model.solve(solver)

    import print_ans

    # 收集每年、每块地、每种作物的种植情况
    for t in years:
        for i in lands:
            for j in crops:
                area = x[i, j, t].value()  # 获取决策变量的值
                crop_name = mapping.get_crop_name(j)
                if area > 0:  # 只记录种植面积大于0的情况
                    key = f"{t}_第一季"
                    # 确定地块的索引位置
                    land_index = mapping.lands_index[i]
                    crop_index = mapping.crops_index[crop_name]
                    # 将种植面积填入到对应的结构化数组
                    print_ans.seasonal_data[key][land_index][crop_index] = area
